#!/bin/bash

# @brief 参数使用说明
function Usage() {
    echo "---------------------------------------------------------------------------------"
    echo "$0  version 2.0"
    echo "Usage:"
    echo "本脚本用于运行APU DLA模型"
    echo "-------------------------------------"
    echo "  --help|-h      查看指令帮助信息"
    echo "-------------------------------------"
    echo "  --input|-i    <input_dla_path>           [必须], 输入的so模型或者bin模型"
    echo "  --data_dir|-d <datadir_path>             [必须], 输入数据目录,数据后缀为.raw,如果有多个按照字母序排列"
    echo "  --sdk_home <sdk_root_path>               [必须], 使用的sdk根目录"
    echo "  --output|-o  <output_dir>         		 [可选], 输出的模型目录,默认当前目录下的output"
    echo "  --sdk|-v  <sdk_version>    				 [必须], 使用的NeuroPilot sdk 版本,默认:从sdk目录中提取"
    echo "  --input_num |-n <count>					 [可选], 模型的输入个数,默认:1"
    echo "  --profile | -p        					 [可选], 开启profile ,会在logcat日志中显示模型运行耗时,日志关键字:neuron"
    echo "  logcat tag is: neuron"
    echo "  etc:"
    echo "      ./apu_netrun.sh -s /opt/neuropilot-sdk/neuropilot-sdk-basic-7.0.2-build20231029  -d datas/rgb_1440 -i /data/code/new_project/mt6897/to_GNet_deploy_model_simple_half_tcm_opt3_opt-bw.dla -p"
    echo "---------------------------------------------------------------------------------"
    exit 0
}

##< @brief  日志函数接口
##< 不同日志函数接口，会增加不同的日志输出前缀，如日志级别、文件名、行数、所在函数等
##<
##< @return 成功返回0，失败返回其他
LogError() {
    local prefix_switch=$1
    if [[ "${prefix_switch}" == false ]]; then
        shift
        echo -e "\033[41;30m$*\033[0m" >&2
    else
        local debug_info=$(caller 0)
        local line_num=$(echo ${debug_info} | awk '{print $1}')
        local filename=$(echo ${debug_info} | awk '{print $3}')
        local function_name=$(echo ${debug_info} | awk '{print $2}')
        echo -e "\033[41;30m[E][$(date +'%F %T')][${filename}:${line_num}][${function_name}]: $*\033[0m" >&2
    fi
}

LogInfo() {
    local prefix_switch=$1
    if [[ "${prefix_switch}" == false ]]; then
        shift
        echo -e "$*" >&2
    else
        local debug_info=$(caller 0)
        local line_num=$(echo ${debug_info} | awk '{print $1}')
        local filename=$(echo ${debug_info} | awk '{print $3}')
        local function_name=$(echo ${debug_info} | awk '{print $2}')
        echo -e "[I][$(date +'%F %T')][${filename}:${line_num}][${function_name}]: $*" >&2
    fi
}

LogVerbose() {
    echo -e "\033[34m [V]: $*\033[0m"
}

declare -A sdk_ver_list=(["5.0.2"]="neuropilot-sdk-basic-5.0.2-build20220217" \
    ["5.0.3"]="neuropilot-sdk-basic-5.0.3-build20220318" \
    ["5.0.4"]="neuropilot-sdk-basic-5.0.4-build20220414" \
    ["6.0.0-beta2"]="neuropilot-sdk-basic-6.0.0-beta2-build20220802" \
    ["6.0.1"]="neuropilot-sdk-basic-6.0.1-build20221013" \
    ["6.0.3"]="neuropilot-sdk-basic-6.0.3-build20230118" \
    ["7.0.2"]="neuropilot-sdk-basic-7.0.2-build20231029" \
	["7.0.8"]="neuropilot-sdk-premium-7.0.8-build20240807" \
    ["8.0.1"]="neuropilot-sdk-premium-8.0.1-build20240703" \
    ["8.0.2"]="neuropilot-sdk-premium-8.0.2-build20240731" \
    ["8.0.4"]="neuropilot-sdk-premium-8.0.4-build20240925" \
    ["8.0.9"]="neuropilot-sdk-premium-8.0.9-build20250611")

##< @brief 解析本脚本的命令行输入参数, 并放置在全局变量
##< @param [in]  命令行参数
##< @param [out] 全局变量放在函数首
##< @return 成功返回0，失败返回非零值
function ParseArgs() {
    # 解析参数
    g_input_model=""
    g_output_dir=""
    g_sdk_dir="/opt/neuropilot-sdks"
    g_sdk_version=""
    g_model_name=""
    g_model_input_num=1
    g_data_dir=""
    g_open_profile=false
    g_profile_option=""

    while [ -n "$1" ]; do
        case "$1" in
        --input | -i)
            g_input_model=$2
            shift
            ;;
        --output | -o)
            g_output_dir=$2
            shift
            ;;
        --sdk_home)
            g_sdk_dir=$2
            shift
            ;;
        --sdk | -v)
            g_sdk_version=$2
            shift
            ;;
        --profile | -p)
            g_open_profile=true
            g_profile_option=" -c 100"
            ;;
        --input_num | -n)
            g_model_input_num=$2
            shift
            ;;
        --data_dir | -d)
            g_data_dir=$2
            shift
            ;;
        --help | -h)
            Usage
            exit 0
            ;;
        *)
            LogError "unsupport:$1"
            Usage
            return 1
            ;;
        esac
        shift
    done
    LogInfo "g_input_model:${g_input_model}"

    if [[ ! -f "${g_input_model}" ]]; then
        LogError "g_input_model:${g_input_model} is not exist, while current dir:$(pwd)"
        Usage
        return 1
    fi

    if [[ -z "${g_sdk_dir}" ]] || [[ ! -d "${g_sdk_dir}" ]]; then
        LogError "g_sdk_dir:${g_sdk_dir} is not exist, while current dir:$(pwd)"
        Usage
        return 1
    fi
	
	if [[ -z "${g_sdk_version}" ]];then
       LogError "g_sdk_version is invalid"
	   Usage
	   return 1
    fi
	version_path=${sdk_ver_list[$g_sdk_version]}
    if [ ! -d "${sdk_home}/${version_path}" ];then
        version_path="v${sdk_version}"
    fi
	g_sdk_dir=${g_sdk_dir}/${version_path}
	if [[ ! -d "${g_sdk_dir}" ]]; then
        LogError "the sdk version ${g_sdk_version} path: ${g_sdk_dir} is not exist, while current dir:$(pwd)"
        Usage
        return 1
    fi

    if [[ -z "${g_data_dir}" ]] || [[ ! -d "${g_data_dir}" ]]; then
        LogError "g_data_dir:${g_data_dir} not exist!"
        return 1
    fi    

    g_model_type="${g_input_model##*.}"
	g_model_name=$(basename ${g_input_model})
	
	if [[ "${g_model_type}" != "dla" ]]; then            
		LogError "[Error] only support DLA model file!!"
		Usage
        return 1
	fi

    LogInfo "############### input args: ########################"
    LogInfo "#  input_model      : ${g_input_model}"
    LogInfo "#  output_dir       : ${g_output_dir}"
    LogInfo "#  sdk_dir          : ${g_sdk_dir}"
    LogInfo "#  sdk_version      : ${g_sdk_version}"
    LogInfo "#  model_name       : ${g_model_name}"
    LogInfo "#  model_input_num  : ${g_model_input_num}"
    LogInfo "####################################################"
    return 0
}

# 1.导入脚本基础组件
g_tools_dir=$(dirname "$0")
. ${g_tools_dir}/adb_utils.sh
adb_device
echo "target device ANDROID_VERSION: $ANDROID_VERSION"

# 2. 解析命令行入参数
ParseArgs $@
ret=$?
if [[ $ret -ne 0 ]]; then    
    exit 1
fi
LogInfo "==>ParseArgs OK"

# 3. 推送依赖的库及模型文件
target_dir=/data/local/tmp/ApuNetRun
adb_run mkdir -p ${target_dir}/out
adb_run mkdir -p ${target_dir}/input
adb_run rm -rf ${target_dir}/out/*

target_soc_type=$(adb_getprop ro.soc.model)
LogVerbose "TARGET SOC:${target_soc_type}"
LogVerbose "================================Push libs ==============================>"

target_soc_type=${target_soc_type,,}

soc_dir=${g_sdk_dir}/neuron_sdk/${target_soc_type}
if [ ! -d "${soc_dir}" ];then
	LogVerbose "[ERROR] unsupport soc ${target_soc_type} for sdk ${g_sdk_version}"
    exit 2
fi

LibList=("bin/neuronrt")

listAll=(${LibList[*]})
for lib in ${listAll[@]}; do
    LogVerbose "adb push ${lib}  ${target_dir}"
    adb_push "${soc_dir}/${lib}" "${target_dir}"
    ret=$?
    if [[ $ret -ne 0 ]]; then
        LogVerbose "[ERROR]"
        exit 1
    fi
done

# 2. 推送配置文件及待转换的库

adb_push ${g_data_dir}/*.raw ${target_dir}/input/
dataList=($(ls ${g_data_dir}))
for dataFile in ${dataList[@]}; do
    LogVerbose ":$dataFile"
done

if [[ ${#dataList[@]} -lt ${g_model_input_num} ]];then
    LogError "dataList(${#dataList[@]}) < g_model_input_num(${g_model_input_num})"
    exit 1
fi
input_data_option=""
for ((i=0;i<${g_model_input_num};i++))
do
    input_data_option="${input_data_option} -i input/${dataList[i]}"
done

adb_push ${g_input_model} ${target_dir}
ret=$?
if [[ $ret -ne 0 ]]; then
    LogError "adb push ${g_input_model}  ${target_dir}"
    exit 1
fi

# 3. 远程运行
model_name=${g_input_model##*/}
# adb shell "chmod  +x  ${target_dir}/${remote_script_name}"
echo "==> show dla I/O info:"

if [[ "${g_open_profile}" == true ]]; then
    ProfileCMD="export MTKNN_ENABLE_PROFILER=1 && export MTKNN_ENABLE_PROFILER_DETAIL=1"
else
    ProfileCMD="export MTKNN_ENABLE_PROFILER=0"
fi

RemoteCMD="${ProfileCMD} && cd ${target_dir} && ./neuronrt -m hw -a ${model_name} ${input_data_option} -o out/out.raw -d ${g_profile_option}"
echo "===>"
echo "===> cmd: ${RemoteCMD}"
adb_run "${RemoteCMD}"
ret=$?
if [[ $ret -ne 0 ]]; then
    echo "[ERROR] fail to run: adb shell ${RemoteCMD}"
    exit $ret
else
    echo "neuronrt run net OK!"
fi

adb_pull ${target_dir}/out/out.raw .

if [[ -n "${g_output_dir}" ]];then
    mkdir -p ${g_output_dir}
    mv out.raw  ${g_output_dir}
fi
echo "Done."